import tensorflow as tf
import matplotlib.pyplot as plt

# 构造Feature结构，告诉解码器每个Feature是什么
feature_description = {
    'dsm': tf.io.FixedLenFeature([512, 512, 1], tf.float32),
    'rgb': tf.io.FixedLenFeature([], tf.string),
    'label': tf.io.FixedLenFeature([], tf.string)
}


# Example的解析函数
def parse_example(example_string):
    feature = tf.io.parse_single_example(serialized=example_string, features=feature_description)

    feature['rgb'] = tf.image.decode_png(feature['rgb'], channels=3)
    feature['rgb'] = tf.image.resize(feature['rgb'], [512, 512])
    feature['rgb'] = tf.cast(feature['rgb'], tf.float32)
    feature['rgb'] = feature['rgb'] / 255

    feature['label'] = tf.image.decode_png(feature['label'], channels=1)
    feature['label'] = tf.image.resize(feature['label'], [512, 512])
    feature['label'] = tf.cast(feature['label'], tf.int64)
    feature['label'] = tf.squeeze(feature['label'])

    return feature['dsm'], feature['rgb'], feature['label']


if __name__ == '__main__':
    dataset = tf.data.TFRecordDataset(r'Potsdam.tfrecords')
    dataset = dataset.map(parse_example)
    dataset = dataset.shuffle(buffer_size=500)

    for dsm, rgb, label in dataset.take(1):
        temp = (dsm - tf.reduce_min(input_tensor=dsm, axis=(0, 1, 2))) / (
                tf.reduce_max(input_tensor=dsm, axis=(0, 1, 2)) - tf.reduce_min(input_tensor=dsm, axis=(0, 1, 2)))

        plt.figure(figsize=(15, 5), dpi=150)
        plt.subplot(1, 3, 1)
        plt.imshow(temp)
        plt.axis('off')

        plt.subplot(1, 3, 2)
        plt.imshow(rgb)
        plt.axis('off')

        plt.subplot(1, 3, 3)
        plt.imshow(label)
        plt.axis('off')

        plt.show()
